package com.hadwinling.alogriithm.projectforpso.psohomework

import breeze.linalg.sum
import breeze.numerics.pow
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.rdd.RDD
import org.apache.spark.{SparkConf, SparkContext}

import scala.util.Random

object homework {

  val pop_size = 5
  val iter_max = 200
  val dim = 2
  var t = 0
  val c1 = 2
  val c2 = 2
  val w = 1
  //初始化全局最优解（fitness,pos,vel）
  var gb: (Double, (Array[Double], Array[Double])) = (Double.PositiveInfinity, (Array.fill(dim)(0.0), Array.fill(dim)(0.0)))

  def main(args: Array[String]): Unit = {
    val partitionsNum = 4 //设置电脑cpu的核心数量
    val conf = new SparkConf().setAppName("SparkPSO").setMaster("local")

    val sc = new SparkContext(conf)
    //初始化每个粒子(num,(pos,vel,pb_pos))
    val particals_tmp: RDD[(Int, (Array[Double], Array[Double], Array[Double]))] = sc.parallelize(List.range(0, pop_size), partitionsNum)
      .map(num => (num, (Array.fill(dim)(Random.nextDouble()), Array.fill(dim)(Random.nextDouble()), Array.fill(dim)(Random.nextDouble()))))
    var particles_fitness: RDD[(Int, (Array[Double], Array[Double], Double, Array[Double], Double))] = particals_tmp.mapPartitions(addFitness).persist()
    //  粒子设置为全局广播
    var particlesBC: Broadcast[collection.Map[Int, (Array[Double], Array[Double], Array[Double])]] = sc.broadcast(particles_fitness.map(x => (x._1, (x._2._1, x._2._2, x._2._4))).collectAsMap)
    gb = particles_fitness.map(x => (x._2._3, (x._2._1, x._2._2))).sortByKey(true).take(1)(0)
    //gbBC:(fitness,(pos,vel))
    //找到全局最优解
    var gbBC: Broadcast[(Double, (Array[Double], Array[Double]))] = sc.broadcast(gb)
    t = t + 1
    /* calling the function and update the information of particles */
    var particlesArray: Array[(Int, (Array[Double], Array[Double], Array[Double]))] = sc.parallelize(List.range(0, pop_size), partitionsNum).map { x =>
      //      更新粒子
      updateParticle(x, particlesBC.value(x)._1, particlesBC.value(x)._2, particlesBC.value(x)._3, gbBC.value._2._1)
    }.collect
    //cleaning
    particlesBC.destroy

    while (t < iter_max) {
      particles_fitness = sc.parallelize(particlesArray, partitionsNum).mapPartitions(addFitness).persist()
      particlesBC = sc.broadcast(particles_fitness.map(x => (x._1, (x._2._1, x._2._2, x._2._4))).collectAsMap)
      /*find the global best */
      gb = particles_fitness.map(x => (x._2._3, (x._2._1, x._2._2))).sortByKey(true).take(1)(0)
      gbBC.destroy
      gbBC = sc.broadcast(gb)
      List.range(0, pop_size).foreach { j =>
        println("num:" + particlesArray(j)._1)
        print("gb_pos:")
        List.range(0, dim).foreach { i =>
          print(gb._2._1(i) + "\t")
        }
        println
        println("gb_fitness:" + gb._1)
        print("pos:")
        List.range(0, dim).foreach { i =>
          print(particlesArray(j)._2._1(i) + "\t")
        }
        println
        print("vel:")
        List.range(0, dim).foreach { i =>
          print(particlesArray(j)._2._2(i) + "\t")
        }
        println
        print("pb_pos:")
        List.range(0, dim).foreach { i =>
          print(particlesArray(j)._2._3(i) + "\t")
        }
        println
        println
      }
      t = t + 1
      particlesArray = sc.parallelize(List.range(0, pop_size), partitionsNum).map { x =>
        updateParticle(x, particlesBC.value(x)._1, particlesBC.value(x)._2, particlesBC.value(x)._3, gbBC.value._2._1)
      }.collect

      //cleaning
      particlesBC.destroy
    }
    sc.stop()
  }

  def myFun(pos: Array[Double]) = {
    sum(pos.map(pos => pow(pos, 2)))
  }

  //  (num,(pos,vel,pb_pos))
  def addFitness(particle: Iterator[(Int, (Array[Double], Array[Double], Array[Double]))]) = {
    particle.map {
      x =>
        val fitness = myFun(x._2._1)
        val pb_fitness = myFun(x._2._3)
        //  (num,(pos,vel,fitness,pb_pos,pb_fitness))
        (x._1, (x._2._1, x._2._2, fitness, x._2._3, pb_fitness))
    }
  }

  //更新粒子操作
  def updateParticle(num: Int, pos: Array[Double], vel: Array[Double], pb_pos: Array[Double], gb_pos: Array[Double]) = {
    val newpos: Array[Double] = Array.fill(dim)(0.0)
    val newvel: Array[Double] = Array.fill(dim)(0.0)
    val newpb_pos: Array[Double] = Array.fill(dim)(0.0)
    val fitness: Double = myFun(pos)
    var newfitness: Double = 0.0
    List.range(0, dim).foreach { i =>
      newvel(i) = w * vel(i) + c1 * Random.nextDouble() * (pb_pos(i) - pos(i)) + c2 * Random.nextDouble() * (gb_pos(i) - pos(i))
      newpos(i) = pos(i) + newvel(i)
    }
    newfitness = myFun(newpos)
    if (newfitness < fitness) {
      List.range(0, dim).foreach { i =>
        newpb_pos(i) = newpos(i)
      }
    } else {
      List.range(0, dim).foreach { i =>
        newpb_pos(i) = pb_pos(i)
      }
    }
    (num, (newpos, newvel, newpb_pos))
  }

}
